import argparse
import os

import torch
# import wandb
import yaml

from Aggregation.AdapterClip import AdapterClip
from Aggregation.FedAvg import FedAvg
from Aggregation.FedNova import FedNova
from Aggregation.FedProx import FedProx
from Aggregation.Generation import Generation
from Aggregation.Scaffold import Scaffold
from Aggregation.FedGPMAEfficient2 import FedGPMAEfficient2
from data.datasets import partition_data
from utils.Seed import setup_seed
from utils.options import args_parser


def load_config(args):
    with open(args.config) as file:
        try:
            databaseConfig = yaml.safe_load(file)
            return databaseConfig
        except yaml.YAMLError as exc:
            print(exc)


if __name__ == '__main__':
    os.environ["TOKENIZERS_PARALLELISM"] = "false"

    config = args_parser()
    # config = load_config(args)
    # wandb.init(config=config, project="FedGXNU")

    setup_seed(config.seed)
    train_dataset, test_dataset, user_groups, traindata_cls_counts = partition_data(
        config.dataset, config.partition, beta=config.beta, num_users=config.num_users)

    # device = config.device.
    torch.cuda.set_device(config.device)

    aggregation = None
    config.aggregation = config.aggregation.lower().strip()
    if config.aggregation == "fedavg":
        aggregation = FedAvg(config, train_dataset, test_dataset, user_groups, traindata_cls_counts)
    elif config.aggregation == "fedprox":
        aggregation = FedProx(config, train_dataset, test_dataset, user_groups, traindata_cls_counts)
    elif config.aggregation == "fednova":
        aggregation = FedNova(config, train_dataset, test_dataset, user_groups, traindata_cls_counts)
    elif config.aggregation == "scaffold":
        aggregation = Scaffold(config, train_dataset, test_dataset, user_groups, traindata_cls_counts)
    elif config.aggregation == "generator":
        aggregation = Generation(config, train_dataset, test_dataset, user_groups, traindata_cls_counts)
    elif config.aggregation == "fedgpma":
        aggregation = AdapterClip(config, train_dataset, test_dataset, user_groups, traindata_cls_counts)
    elif config.aggregation == "fedgpma-efficient":
        aggregation = FedGPMAEfficient2(config, train_dataset, test_dataset, user_groups, traindata_cls_counts)

    aggregation.train()